# codegen_8086_dos.py

from collections import namedtuple
from middleCode import Quad

class AssemblyGenerator8086DOS:
    def __init__(self):
        self.asm = []
        self.strings = {}
        self.str_count = 0
        self.op_map = {
            '+': 'ADD', '-': 'SUB', '*': 'MUL', '/': 'DIV', '%': 'MOD',
            '>': 'GT', '<': 'LT', '>=': 'GE', '<=': 'LE', '==': 'EQ', '!=': 'NE',
            '&&': 'AND', '||': 'OR'  # 新增逻辑运算符
        }
        self._label_count = 0
        self.consts = {}  # 用于存储常量映射
        self.const_count = 0

    def new_const_label(self, value):
        if value not in self.consts:
            self.const_count += 1
            label = f"CONST{self.const_count}"
            self.consts[value] = label
        return self.consts[value]

    def new_str_label(self, literal):
        if literal not in self.strings:
            self.str_count += 1
            self.strings[literal] = f"STR{self.str_count}"
        return self.strings[literal]

    def new_label(self, base):
        self._label_count += 1
        return f"{base}_{self._label_count}"

    # 辅助方法：获取操作数
    def _get_operand(self, operand):
        if operand is None:
            return None
        if isinstance(operand, str) and operand.isdigit():
            return self.new_const_label(operand)
        return operand
    
    def generate(self, quads):
         # —— 1. 收集所有 VAR_DECL —— 
        decls = [(q.arg1, q.arg2, q.res)
                 for q in quads if q.op == 'VAR_DECL']

        # —— 2. DATA 段头 —— 
        self.asm += [
            "    .MODEL SMALL",
            "    .STACK 100h",
            "",
            "    .DATA",
        ]
        # 字符串常量
        for q in quads:
            if q.op == 'LOAD_STR':
                self.new_str_label(q.arg1)
        for lit, lbl in self.strings.items():
            esc = lit.replace("'", "''")
            self.asm.append(f"{lbl} DB '{esc}','$'")

        # 2. 数字常量声明
        for q in quads:
            # 二元算术/赋值里出现的立即数
            for lit in (q.arg1, q.arg2):
                # 只有当 lit 是字符串且全是数字时，才当作立即数
                if isinstance(lit, str) and lit.isdigit():
                    lbl = self.new_const_label(lit)
        for val, lbl in self.consts.items():
            # 在 DATA 段为每个常量添加一行
            self.asm.append(f"    {lbl} DW {val}")
            
        # —— 3. 由 VAR_DECL 定义变量 —— 
        for name, typ, scope in decls:
            self.asm.append(f"    {name} DW ?    ; {scope} {typ}")

        # 临时变量
        temps = sorted({q.res for q in quads if q.res and q.res.startswith('t')})
        for t in temps:
            self.asm.append(f"    {t} DW ?")

        self.asm += ["", "    .CODE", ""]

        # ==== PRINT_NUMBER 子程序 ====
        self.asm += [
            "PRINT_NUMBER PROC",
            "    PUSH AX",
            "    PUSH BX",
            "    PUSH CX",
            "    PUSH DX",
            "    MOV CX,0",
            "    MOV BX,10",
            "    CMP AX,0",
            "    JNE _PN_LOOP",
            "    MOV DL,'0'",     "    MOV AH,02h", "    INT 21h",
            "    JMP _PN_DONE",
            "_PN_LOOP:",
            "    XOR DX,DX",      "    DIV BX",      "    PUSH DX",
            "    INC CX",         "    CMP AX,0",    "    JNE _PN_LOOP",
            "_PN_OUT:",
            "    POP DX",         "    ADD DL,'0'",  "    MOV AH,02h",
            "    INT 21h",        "    LOOP _PN_OUT",
            "_PN_DONE:",
            "    POP DX",
            "    POP CX",
            "    POP BX",
            "    POP AX",
            "    RET",
            "PRINT_NUMBER ENDP",
            "SCANF_NUMBER PROC",
            "    PUSH BX",
            "    PUSH CX",
            "    PUSH DX",
            "    MOV DX, 0",        
            "    MOV CX, 0",        
            "    MOV BX, 10",      
            "_READ_CHAR:",
            "    MOV AH, 01h",      
            "    INT 21h",
            "    CMP AL, 0DH",      
            "    JE  _END_INPUT",
            "    CMP AL, '0'",      
            "    JL  _END_INPUT",
            "    CMP AL, '9'",
            "    JG  _END_INPUT",
            "    MOV AH, 0",
            "    SUB AL, '0'",      
            "    MOV CX, AX",       
            "    MOV AX, DX",       
            "    MUL BX",           
            "    MOV DX, AX",       
            "    MOV AX, CX",       
            "    ADD DX, AX",       
            "    JMP _READ_CHAR",
            "_END_INPUT:",
            "    MOV AX, DX",       
            "    POP DX",           
            "    POP CX",
            "    POP BX",
            "    RET",
            "SCANF_NUMBER ENDP",
            ""
        ]
        # 用户函数（除 main），带标签映射
        current = None
        label_map = {}
        # 计算每个函数的局部大小
        locals_map = {}
        for q in quads:
            if q.op == 'FUNC_BEGIN':
                current = q.arg1; locals_map[current] = 0
            elif q.op == 'VAR_DECL' and current:
                if q.res == 'local': locals_map[current] += 2
            elif q.op == 'FUNC_END': current = None
        # 生成函数体
        current = None; label_map = {}
        for q in quads:
            if q.op == 'FUNC_BEGIN':
                current = q.arg1
                if current != 'main':
                    size = locals_map.get(current, 0)
                    self.asm += [f"{current} PROC", "    PUSH BP", "    MOV BP, SP"]
                    if size: self.asm.append(f"    SUB SP, {size}    ; alloc locals")
                    label_map = {}
                continue
            if q.op == 'FUNC_END':
                if current != 'main':
                    self.asm += ["    MOV SP, BP", "    POP BP", "    RET", f"{current} ENDP"]
                current = None
                label_map = {}
                continue
            # 重映射静态标签
            if current and current != 'main':
                if q.op == 'LABEL':
                    new = label_map.setdefault(q.res, f"{current}_{q.res}")
                    self.asm.append(f"{new}:")
                    continue
                if q.op in ('JMP','JZ','JNZ'):
                    target = q.res
                    newt = label_map.setdefault(target, f"{current}_{target}")
                    if q.op == 'JMP': self.asm.append(f"    JMP {newt}")
                    else:
                        # 复制原 emit_JZ/JNZ 逻辑，但用 newt
                        cond = 'JE SHORT' if q.op=='JNZ' else 'JNE SHORT'
                        skip = self.new_label(f"{current}_SKIP")
                        self.asm += [f"    MOV AX, {q.arg1}", "    CMP AX, 0",
                                     f"    {cond} {skip}", f"    JMP {newt}", f"{skip}:"]
                    continue
                # 正常四元式
                if q.op == 'RET':
                    self.asm += [f"    MOV AX, {q.arg1}", "    MOV SP, BP", "    POP BP", "    RET"]
                    continue
                if q.op == 'ARG':
                    self.asm.append(f"    PUSH {q.arg1}")
                    continue
                if q.op == 'CALL': self._emit_call_generic(q); continue
                # 其余用默认
                self._emit_by_op(q)
                continue
        # MAIN 部分
        self.asm += ["", "MAIN PROC", "    MOV AX,@DATA", "    MOV DS,AX", ""]
        for q in quads:
            if q.op in ('FUNC_BEGIN','FUNC_END'): continue
            if q.op == 'ARG': self.asm.append(f"    PUSH {q.arg1}"); continue
            if q.op == 'CALL': self._emit_call_generic(q); continue
            if q.op == 'RET': continue
            self._emit_by_op(q)
        self.asm += ["", "    MOV AH,4Ch", "    INT 21h", "MAIN ENDP", "    END MAIN"]
        return "\n".join(self.asm)
    
    def _emit_by_op(self, q):
        # 映射到 emit_xxx
        op = self.op_map.get(q.op, q.op)
        handler = getattr(self, f"emit_{op}", self.emit_unknown)
        handler(q)

    def _emit_call_generic(self, q):
        if q.arg1 in ('printf','scanf'): self._process_call(q)
        else:
            self.asm.append(f"    CALL {q.arg1}")
            cnt = q.arg2 or 0
            if cnt: self.asm.append(f"    ADD SP, {cnt*2}    ; clean args")
            if q.res: self.asm.append(f"    MOV {q.res}, AX")

    def emit_FUNC_BEGIN(self, q):
        name = q.arg1
        self.asm += [f"{name} PROC", "    PUSH BP", "    MOV BP, SP"]
        # TODO: 如果有局部变量，SUB SP, total_locals_size

    def emit_RET(self, q):
        # 假设返回值放在 q.arg1
        self.asm += [f"    MOV AX, {q.arg1}", "    MOV SP, BP", "    POP BP", "    RET"]

    def emit_FUNC_END(self, q):
        name = q.arg1
        self.asm.append(f"{name} ENDP")

    def emit_ARG_for_call(self, q):
        self.asm.append(f"    PUSH {q.arg1}")
    
    # 专门处理 CALL
    def emit_CALL(self, q):
        func = q.arg1
        if func in ('printf','scanf'):
            self._process_call(q)
        else:
            # 假设都是 cdecl：push  参数已在 ARG 阶段生成
            self.asm.append(f"    CALL {func}")
            # 参数个数在 q.arg2 或者其他字段可获知
            n = q.arg2 or 0
            if n > 0:
                self.asm.append(f"    ADD SP, {n*2}  ; 清理 {n} 个参数")
            # 把返回值放到 res
            if q.res:
                self.asm.append(f"    MOV {q.res}, AX")
    # ---------- 赋值 & 算术 ----------
    def emit_ASSIGN(self, q):
        # 确保源操作数是立即数/变量，目标是变量
        if q.arg1.isdigit():  # 处理立即数赋值
            self.asm += [
                f"    MOV AX, {q.arg1}",
                f"    MOV {q.res}, AX"
            ]
        else:  # 处理变量赋值
            self.asm += [
                f"    MOV AX, {q.arg1}",
                f"    MOV {q.res}, AX"
            ]

    def emit_ADD(self, q):
        arg1 = self._get_operand(q.arg1)
        arg2 = self._get_operand(q.arg2)
        self.asm += [
            f"    MOV AX, {arg1}",
            f"    ADD AX, {arg2}",
            f"    MOV {q.res}, AX"
        ]

    def emit_SUB(self, q):
        arg1 = self._get_operand(q.arg1)
        arg2 = self._get_operand(q.arg2)
        self.asm += [
            f"    MOV AX, {arg1}",
            f"    SUB AX, {arg2}",
            f"    MOV {q.res}, AX"
        ]

    def emit_MUL(self, q):
        arg1 = self._get_operand(q.arg1)
        arg2 = self._get_operand(q.arg2)
        self.asm += [
            f"    MOV AX, {arg1}",
            f"    MUL {arg2}",
            f"    MOV {q.res}, AX"
        ]

    def emit_DIV(self, q):
        arg1 = self._get_operand(q.arg1)
        arg2 = self._get_operand(q.arg2)
        self.asm += [
            f"    MOV AX, {arg1}",
            "    XOR DX, DX",
            f"    DIV {arg2}",
            f"    MOV {q.res}, AX"
        ]

    def emit_MOD(self, q):
        arg1 = self._get_operand(q.arg1)
        arg2 = self._get_operand(q.arg2)
        self.asm += [
            f"    MOV AX, {arg1}",
            "    XOR DX, DX",
            f"    DIV {arg2}",
            f"    MOV {q.res}, DX"
        ]


    # ---------- 关系运算 ----------
    def _emit_relop(self, q, jmp_instr):
        left, right, res = q.arg1, q.arg2, q.res
        L_true = self.new_label("REL_TRUE")
        L_end  = self.new_label("REL_END")
        self.asm += [
            f"    MOV AX,{left}",
            f"    CMP AX,{right}",
            f"    MOV {res},0",
            f"    {jmp_instr} {L_true}",
            f"    JMP {L_end}",
            f"{L_true}:",
            f"    MOV {res},1",
            f"{L_end}:"
        ]

    # 新增逻辑运算处理
    def emit_AND(self, q):
        arg1, arg2, res = q.arg1, q.arg2, q.res
        L_false = self.new_label("AND_FALSE")
        L_end = self.new_label("AND_END")
        self.asm += [
            f"    MOV AX, {arg1}", "    CMP AX, 0", f"    JE {L_false}",
            f"    MOV AX, {arg2}", "    CMP AX, 0", f"    JE {L_false}",
            f"    MOV {res}, 1", f"    JMP {L_end}",
            f"{L_false}:", f"    MOV {res}, 0",
            f"{L_end}:"
        ]

    def emit_OR(self, q):
        arg1, arg2, res = q.arg1, q.arg2, q.res
        L_true = self.new_label("OR_TRUE")
        L_end = self.new_label("OR_END")
        self.asm += [
            f"    MOV AX, {arg1}", "    CMP AX, 0", f"    JNE {L_true}",
            f"    MOV AX, {arg2}", "    CMP AX, 0", f"    JNE {L_true}",
            f"    MOV {res}, 0", f"    JMP {L_end}",
            f"{L_true}:", f"    MOV {res}, 1",
            f"{L_end}:"
        ]

    def emit_GT(self, q): self._emit_relop(q, "JG")
    def emit_LT(self, q): self._emit_relop(q, "JL")
    def emit_GE(self, q): self._emit_relop(q, "JGE")
    def emit_LE(self, q): self._emit_relop(q, "JLE")
    def emit_EQ(self, q): self._emit_relop(q, "JE")
    def emit_NE(self, q): self._emit_relop(q, "JNE")

    # ---------- 控制流 ----------
    def emit_JZ(self, q):
        """
        8086 条件跳转只能用 SHORT（±128B），超过后必须用短跳 + 无条件近跳的组合。
        """
        # 为 Skip 生成一个临时标签
        L_skip = self.new_label("JZ_SKIP")
        self.asm += [
            f"    MOV AX, {q.arg1}",
            "    CMP AX, 0",
            # 如果不等于 0，就短跳过下面那条近跳
            f"    JNE SHORT {L_skip}",
            # 等于 0 时，走到这里，执行无条件的近跳（16-bit 位移）
            f"    JMP {q.res}",
            # Skip 标签，继续下面的指令流
            f"{L_skip}:"
        ]

    def emit_JNZ(self, q):
        L_skip = self.new_label("JNZ_SKIP")
        self.asm += [
            f"    MOV AX, {q.arg1}",
            "    CMP AX, 0",
            f"    JE SHORT {L_skip}",
            f"    JMP {q.res}",
            f"{L_skip}:"
        ]
    
    def emit_JMP(self, q):
        self.asm.append(f"    JMP {q.res}")  # 近跳转

    def emit_LABEL(self, q):
        self.asm.append(f"{q.res}:")

    # ---------- I/O & CALL ----------
    def emit_LOAD_STR(self, q):
        lbl = self.new_str_label(q.arg1)
        self.asm += [
            f"    LEA DX,{lbl}",
            "    MOV AH,09h",
            "    INT 21h"
        ]

    def emit_ARG(self, q):
        # 对非字符串参数，调用 PRINT_NUMBER
        if not q.arg1.startswith("STR"):
            self.asm += [
                f"    MOV AX,{q.arg1}",
                "    CALL PRINT_NUMBER"
            ]

    def _process_call(self, q):
        func = q.arg1           # 'printf' 或 'scanf'
        var  = q.res            # 对于 printf，用来打印的值；对于 scanf，用来存放的变量名
        # 从调用四元式里再追溯一下参数节点的类型（你可能需要在 AST 或 Quad 里带上原始节点类型）
        param_type = q.param_type

        if func == 'printf':
            if param_type == 'string':
                # 输出字符串
                self.asm += [
                    f"    LEA DX, {var}",    # var 是串的标签
                    "    MOV AH, 09h",
                    "    INT 21h"
                ]
            else:
                # 数字走 PRINT_NUMBER
                self.asm += [
                    f"    MOV AX, {var}",
                    "    CALL PRINT_NUMBER"
                ]

        elif func == 'scanf':
            if param_type == 'string':
                # 调用一个专门的读取字符串例程
                self.asm += [
                    f"    LEA DX, {var}",
                    "    CALL SCANF_STRING"   # 你需要自己实现这个例程
                ]
            else:
                # 数字走原 SCANF_NUMBER
                self.asm += [
                    "    CALL SCANF_NUMBER",
                    f"    MOV {var}, AX"
                ]

    def emit_unknown(self, q):
        # 不再输出“未实现”的警告
        # self.asm.append(f"    ; 未实现的操作: {q.op} {q.arg1} {q.arg2} {q.res}")
        return 

def run_object_code(quads):
    gen = AssemblyGenerator8086DOS()
    return gen.generate(quads)
